import json
from pocketflow import AsyncNode 
from mcp.types import Tool, TextContent 

from typing import Dict, Tuple, List, Any
from fastmcp.exceptions import ToolError
from scorpio.core.mcp import MCPClient
from scorpio.core.schemas import Message, MessageToolCall
from .utils import parse_json

from scorpio.core.common import get_logger
logger = get_logger(__name__)

"""
shared should include:
    - MCP and LLM session
    - Call tools from LLM
    - others? need to explore...
    - previous node, previous inputs?
"""


class MCPAction(AsyncNode):
    async def prep_async(self, shared):
        """Prepare the context and question for mcp action process."""
        # Get the current context (default to "No previous search" if none exists)
        context = shared.get("context")
        messages=context.get("messages")
        return messages, context
        
    async def exec_async(self, inputs):
        """Call the LLM to decide whether to search or answer."""
        messages, context = inputs
        message:Message=messages[-1]
        # check complete in prep phase, return complete at here
        cli: MCPClient= context.get("mcp_client")
        tool_calls: List[MessageToolCall]=message.tool_calls
        logger.info(f"current action: call tool {[tc.function.name for tc in tool_calls]} with message id:{message.id}")
        content:List[Dict]=[]
        try:
            for idx, tc in enumerate(tool_calls):
                arguments=parse_json(tc.function.arguments)
                result = await cli.client.call_tool(tc.function.name, arguments)
                tool_message=Message(
                    role="tool",
                    name=tc.function.name,
                    content=result.content[0].text)
                if tc.id is not None:
                    tool_message.tool_call_id=tc.id
                messages.append(tool_message)
                content.append(tool_message)
                logger.info(f"tool call index {idx}: call result\n {tool_message.content}")
            # to-do: record event call tool
            #print(type(result[0]))
            #messages.extend()
        except ToolError as e:
            return {
                "action": "failure",
                "text": f"mcp call_tool error: {e}"
            }
        return {
            "action":"success",
            "content": content,
        }
 
    # post_async: return key(str) which used as condition in flow
    async def post_async(self, shared, prep_res, exec_res)->str:
        """
        function:
            - metrics
            - count token, step/node into shared
            - timeout control
        # to-do: package message and conversation
        """
        #metrics=shared["context"].get("metrics")
        #nodes=metrics.get("nodes")
        #nodes.append(self.__class__.__name__)

        if exec_res["action"] == "success":
            # to-do: package message and conversation 
            return "success"
        else:
            return "failure" 
        